from IPython.kernel import client

from anycloud import config
from anycloud.task import TASK_SUCCESS, TASK_SUBMITTED
from anycloud.engine.base import Engine
from anycloud.error import NotFoundError


payload_expression = """
import os, base64

x = payload
file('__payload__.tar', 'wb').write(base64.standard_b64decode(payload))
os.system('tar -xf __payload__.tar')
%s
"""

class IPythonEngine(Engine):
    _params = (
        config.StringParameter(
            'ipython.furl', 'FURL for IPython controller',
            default='',
        ),
    )

    def __init__(self, **kwargs):
        super(IPythonEngine, self).__init__(**kwargs)
        self.tc = client.get_task_client(self.furl)

    def _ipytask(self, task):
        if task.files or task.modules:
            payload = task._export_payload(as_string=True)
            task.push.update({'payload':payload})
            task.expression = payload_expression % task.expression

        return client.StringTask(task.expression, task.pull, task.push, True, True) 

    def submit(self, tasks):
        for task in tasks:
            task.engineinfo['id'] = self.tc.run(self._ipytask(task))
            task.status = TASK_SUBMITTED
        
        return tasks

    def retrieve(self, tasks):
        taskids = [t.engineinfo['id'] for t in tasks]
        self.tc.barrier(taskids)

        for task in tasks:
            task.result = self.tc.get_task_result(task.engineinfo['id']).results
            task.status = TASK_SUCCESS
            # TODO: task.error

        return [t.result for t in tasks]

        
